package top.codeora.aiollama.agent;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import lombok.Data;
import lombok.EqualsAndHashCode;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.tool.ToolCallback;
import top.codeora.aiollama.agent.model.AgentState;

import java.util.List;
import java.util.stream.Collectors;

@EqualsAndHashCode(callSuper = true)
@Data
@Slf4j
public class ToolCallAgent extends ReActAgent {


    //可用工具
    private final ToolCallback[] availableTools;

    //保存工具调用信息的响应结果
    private ChatResponse toolCallResponse;

    //工具调用管理
    private final ToolCallingManager toolCallingManager;

    //禁用Spring AI内置的工具调用机制，自己维护选项和上下文
    private final ChatOptions chatOptions;

    public ToolCallAgent(ToolCallback[] availableTools) {
        super();
        this.availableTools = availableTools;
        this.toolCallingManager = ToolCallingManager.builder().build();
        this.chatOptions = ToolCallingChatOptions.builder().internalToolExecutionEnabled(false).build();

    }

    @Override
    public boolean think() {
        try {
            //1.校验提示词，拼接用户提示词
            if (StrUtil.isNotBlank(getNextStepPrompt())) {
                UserMessage userMessage = new UserMessage(getNextStepPrompt());
                getMessageList().add(userMessage);
            }
            //2.调用 AI大模型，获取工具调用结果
            List<Message> mesageList = getMessageList();
            Prompt prompt = new Prompt(mesageList, this.chatOptions);
            ChatResponse chatResponse = getChatClient().prompt(prompt)
                    .system(getSystemPrompt())
                    .tools(availableTools)
                    .call()
                    .chatResponse();
            this.toolCallResponse = chatResponse;
            //3.解析工具调用结果，获取要调用的工具
            AssistantMessage assistantMessage = chatResponse.getResult().getOutput();
            //获取所有工具调用列表
            List<AssistantMessage.ToolCall> toolCallList = assistantMessage.getToolCalls();
            String result = assistantMessage.getText();
            log.info(getName() + "的思考:{}", result);
            log.info("{}选择了{}个工具", getName(), toolCallList.size());

            String toolCallInfo = toolCallList.stream()
                    .map(toolcall -> String.format("工具名称：%s,参数：%s", toolcall.name(), toolcall.arguments())).collect(Collectors.joining("\n"));
            log.info(toolCallInfo);

            if (toolCallList.isEmpty()) {
                getMessageList().add(assistantMessage);
                return false;
            } else {
                return true;
            }
        } catch (Exception e) {
            log.error("调用AI大模型失败", e);
            return false;
        }

    }

    /**
     * 执行工具调用并处理结果
     *
     * @return
     */
    @Override
    public String act() {
        if (!toolCallResponse.hasToolCalls()) {
            return "没有工具调用";
        }
        Prompt prompt = new Prompt(getMessageList(), this.chatOptions);
        //调用工具
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, toolCallResponse);
        //记录消息上下文
        setMessageList(toolExecutionResult.conversationHistory());
        ToolResponseMessage toolResponseMessage = (ToolResponseMessage) CollUtil.getLast(toolExecutionResult.conversationHistory());

        boolean terminateToolCalled = toolResponseMessage.getResponses().stream().anyMatch(respone -> respone.name().equals("doTerminate"));
        if (terminateToolCalled) {
            //任务结束，更改状态
            setState(AgentState.FINISHED);
        }

        String results = toolResponseMessage.getResponses().stream()
                .map(response -> "工具" + response.name() + "返回的结果" + response.responseData())
                .collect(Collectors.joining("\n"));
        return results;
    }
}
